#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <sys/types.h>

#include <mutex>
#include <thread>

#include <algorithm>
#include <string>
#include <memory>
#include <vector>
#include <map>
#include <set>

using namespace std;

#define now_ms() \
({ \
    struct timeval tv; \
    gettimeofday(&tv, NULL); \
    (tv.tv_sec * 1000 + tv.tv_usec / 1000); \
})

bool g_finish = false;

void handle_signal_intr(int signal)
{
    printf("exit and dump record file\n");
    g_finish = true;
}

struct PidStat
{
    int     pid;
    string  name;
    string  cpu;
    string  rss;
};

struct PidStatLog
{
    int timestamp;
    map<string, PidStat> stats;
};

struct PidInfo
{
    int     pid;
    string  name;
    bool    finish;

    mutex   locker;
    vector<PidStatLog> logs;
};

template <class... Args>
static string format_string(const string& fmt, Args... args)
{
    int size = snprintf(nullptr, 0, fmt.c_str(), args...) + 1;
    auto buf = shared_ptr<char>(new char[size], [](char* p) {delete[] p;});
    snprintf(buf.get(), size, fmt.c_str(), args...);
    return string(buf.get());
}

vector<string> split_string(const string& data, const string& needle)
{
    vector<string> results;
    char* beg = (char*)data.c_str();
    char* pos = NULL;
    while ( ( pos = strstr(beg, needle.c_str()) ) )
    {
        if (beg != pos)
        {
            results.push_back(string(beg, pos - beg));
        }
        beg = pos + needle.length();
    }
    if (beg < data.c_str() + data.length())
    {
        results.push_back(beg);
    }
    return results;
}

void ltrim(std::string &s)
{
    s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](unsigned char ch) {
        return !std::isspace(ch);
    }));
}
 
void rtrim(std::string &s)
{
    s.erase(std::find_if(s.rbegin(), s.rend(), [](unsigned char ch) {
        return !std::isspace(ch);
    }).base(), s.end());
}

void trim(std::string &s)
{
    ltrim(s);
    rtrim(s);
}

string timestamp_to_datetime(time_t timestamp)
{
    struct tm tmTime;

    localtime_r(&timestamp, &tmTime);

    return format_string(
        "%02d:%02d:%02d", 
        tmTime.tm_hour,
        tmTime.tm_min,
        tmTime.tm_sec
    );
}

int get_process_result(const string& command, string& result)
{
    FILE* fp = popen(command.c_str(), "r");
    if (!fp) {
        printf("popen fail, command:%s\n", command.c_str());
        return -1;
    }

    result.clear();
    
    char buf[10240]; 
    while (fgets(buf, sizeof(buf), fp)) {
        result.append(buf, strlen(buf));
    }

    pclose(fp);

    return 0;
}

// PID USER      PR  NI    VIRT    RES    SHR S %CPU %MEM     TIME+ COMMAND
// 9520 root      20   0   39436   1136    932 S  0.0  0.1   0:00.13 main

int get_process_info(int pid, map<string, PidStat>& infos)
{
    auto cmd = format_string("top -Hp %d -b -n 1", pid);
    
    string strResult;
    int ret = get_process_result(cmd, strResult);
    if (ret != 0) {
        return -1;
    }

    bool start_analyse = false;

    auto lines = split_string(strResult, "\n");
    for (auto item : lines) {
        if (item.find("PID") != string::npos && item.find("COMMAND") != string::npos) {
            start_analyse = true;
           // printf("set analyse, item:%s\n", item.c_str());
            continue ;
        }
        if (!start_analyse) {
            continue ;
        }
        trim(item);
        if (item.length() == 0)
        {
            continue ;
        }
        auto arr = split_string(item, " ");
        if (arr.size() != 12) {
            continue ;
        }

        auto pid  = atoi(arr[0].c_str());
        auto name = arr[11]; 
        auto cpu  = arr[8];
        auto rss  = arr[5];

        auto pid_name = format_string("%d_%s", pid, name.c_str());

        infos[pid_name] = PidStat{
            pid,
            name,
            cpu,
            rss
        };
    }

    return 0;
}

/*
line:  1195 /usr/sbin/alsactl -s -n 19 -c -E ALSA_CONFIG_PATH=/etc/alsa/alsactl.conf --initfile=/lib/alsa/init/00main rdaemon
line:  9520 ./main
line: 12172 ./ps_cpu main
line: 12173 sh -c ps -eo pid,cmd | grep main
line: 12175 grep main
*/

int get_pid_names(const string& name, map<int, string>& infos)
{
    string result;
    string cmd = format_string("ps -eo pid,cmd | grep %s", name.c_str());

    int ret = get_process_result(cmd, result);
    if (ret != 0) {
        printf("exec cmd fail, cmd:%s\n", cmd.c_str());
        return -1;
    }

    auto lines = split_string(result, "\n");
    for (auto item : lines) {
        if (item.find("ps_cpu") != string::npos) {
            continue ;
        }
        if (item.find("grep") != string::npos) {
            continue ;
        }
        trim(item);
        auto arr = split_string(item, " ");
        if (arr.size() < 2) {
            continue ;
        }
        infos[atoi(arr[0].c_str())] = arr[1];
    }
    return 0;
}

/*
top - 10:37:11 up 4 days, 10:52,  8 users,  load average: 5.24, 4.85, 4.73
Threads:   4 total,   3 running,   1 sleeping,   0 stopped,   0 zombie
%Cpu(s): 93.3 us,  6.7 sy,  0.0 ni,  0.0 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
KiB Mem :  1863028 total,   272532 free,   992860 used,   597636 buff/cache
KiB Swap:  2097148 total,  1523324 free,   573824 used.   689452 avail Mem 

   PID USER      PR  NI    VIRT    RES    SHR S %CPU %MEM     TIME+ COMMAND                                                          
  9522 root      20   0   39436   1136    932 R 93.8  0.1  53:11.51 worker1                                                          
  9521 root      20   0   39436   1136    932 R 56.2  0.1  53:10.51 worker0                                                          
  9523 root      20   0   39436   1136    932 R 50.0  0.1  52:46.63 worker2                                                          
  9520 root      20   0   39436   1136    932 S  0.0  0.1   0:00.09 main
*/

void pid_monitor(shared_ptr<PidInfo> ptr)
{
    printf("start process monitor, pid:%d\n", ptr->pid);

    while (!ptr->finish) {

        map<string, PidStat> stats;
        int ret = get_process_info(ptr->pid, stats);
        if (ret != 0) {
            break ;
        }

        vector<string> keys;
        for (const auto& item : stats) {
            keys.push_back(item.first);
        }
        sort(keys.begin(), keys.end(), [](const string& a, const string& b) {
            return a < b;
        });

        for (auto& key : keys) {

            auto iter = stats.find(key);
            if (iter == stats.end()) {
                continue ; 
            }
        }

        auto log = PidStatLog{(int)time(0), stats};
        {
            lock_guard<mutex> lock(ptr->locker);
            ptr->logs.push_back(log);
        }

        sleep(1);
    }

    printf("end process monitor, pid:%d\n", ptr->pid);
}

void diff_pid_names(
    const map<int, string>& new_infos, 
    const map<int, string>& old_infos, 
    map<int, string>& add, 
    map<int, string>& del
)
{
    for (const auto& item : new_infos) {
        if (old_infos.find(item.first) == old_infos.end()) {
            add[item.first] = item.second;
        }
    }
    for (const auto& item : old_infos) {
        if (new_infos.find(item.first) == new_infos.end()) {
            del[item.first] = item.second;
        }
    }
}

int dump_file(
    const string& dir, 
    int pid, 
    const string& name, 
    const vector<PidStatLog>& logs
)
{
    string path = format_string("%s/%d_%d.dat", dir.c_str(), pid, time(0));

    FILE* fp = fopen(path.c_str(), "w");
    if (!fp) {
        printf("write_file fail, pid:%d, name:%s, path:%s\n", pid, name.c_str(), path.c_str());
        return -1;
    }

    map<int, string> th_index2name;
    {
        set<string> gather_pid_names;
        for (const auto& log : logs) {
            for (const auto& stat : log.stats) {
                auto iter = gather_pid_names.find(stat.first);
                if (iter == gather_pid_names.end()) {
                    th_index2name[gather_pid_names.size()] = stat.first;
                    gather_pid_names.insert(stat.first);
                }
            }
        }
    }

    string pids_str;
    for (const auto& item : th_index2name) {
        auto field = format_string("%d", item.first);
        pids_str += format_string("%20s", field.c_str());
    }

    do
    {
        auto wfile = [pid, name, fp](const string& data) -> bool {
            string line = format_string("%s\n", data.c_str());
            int ret = fwrite(line.c_str(), 1, line.length(), fp);
            if (ret != line.length()) {
                printf(
                    "write file data fail, pid:%d, name:%s, data.length:%ld\n", 
                    pid, 
                    name.c_str(), 
                    line.length()
                );
                return false;
            }
            return true;
        };

        string data = format_string("pid:%d", pid);
        if (!wfile(data)) {
            break ;
        }

        data = format_string("name:%s", name.c_str());
        if (!wfile(data)) {
            break ;
        }

        if (!wfile(format_string(""))) {
            break ;
        }

        for (auto& item : th_index2name) {
            string line = format_string(
                "thread[%d]:%s", 
                item.first, 
                item.second.c_str()
            );
            if (!wfile(line)) {
                break ;
            }
        }
        if (!wfile("")) {
            break ;
        }

        {
            string field = "t/n";
            if (!wfile(format_string("%10s%s", field.c_str(), pids_str.c_str()))) {
                break ;
            }
        }

        for (auto log : logs) {
            data = "";
            for (int i = 0; i < th_index2name.size(); i++) {
                auto stat = log.stats[th_index2name[i]];
                auto item = format_string(
                    "%s/%s", 
                    stat.cpu.c_str(), 
                    stat.rss.c_str()
                );
                data += format_string("%20s", item.c_str());
            }

            data = format_string(
                "[%s]%20s", 
                timestamp_to_datetime(log.timestamp).c_str(), 
                data.c_str()
            );
            if (!wfile(data)) {
                break ;
            }
        }

    } while (0);

    fclose(fp);

    return 0;
}

int main(int argc, char** argv)
{
    if (argc != 2) {
        printf("usage[0] >> name\n");
        return -1;
    }

    vector<pair<int, string>> cur_pid_infos;

    map<int, string> cur_pid_names;
    map<int, shared_ptr<PidInfo>> pid_info_dt;

    auto monitor_desc = [](const map<int, string>& dt, const string& desc) {
        string pids;
        for (const auto& item : dt) {
            pids += format_string("%d,", item.first);
        }
        if (pids.length() > 0) {
            printf(
                "%s, pids:%s\n", 
                desc.c_str(), 
                pids.substr(0, pids.length() - 1).c_str()
            );
        }
    };

    auto pid_desc = [](const map<int, string>& dt, const string& desc) {
        string pids;
        for (const auto& item : dt) {
            pids += format_string("%d,", item.first);
        }
        if (pids.length() > 0) {
            pids = pids.substr(0, pids.length() - 1);
        }
        printf("%s, pids:%s\n", desc.c_str(), pids.c_str());
    };

    string dir = format_string("./monitor.%ld", time(0));
    mkdir(dir.c_str(), 0644);

    auto dump_files = [dir, &pid_info_dt](const string& desc, const vector<int>& pids = {}) {

        for (const auto& item : pid_info_dt) {

            if (pids.size() > 0 && find(pids.begin(), pids.end(), item.first) == pids.end()) {
                continue ;
            }

            vector<PidStatLog> logs;

            auto ptr = item.second;
            {
                lock_guard<mutex> lock(ptr->locker);
                logs = ptr->logs;
            }

            int ret = dump_file(dir, ptr->pid, ptr->name, logs);
            if (ret != 0) {
                printf(
                    "dump file fail, %s, pid:%d, name:%s\n", 
                    desc.c_str(), 
                    ptr->pid, 
                    ptr->name.c_str()
                );
                return -1;
            }
            printf(
                "dump file succeed, %s, pid:%d, name:%s\n", 
                desc.c_str(), 
                ptr->pid, 
                ptr->name.c_str()
            );
        }
    };

    signal(SIGINT, handle_signal_intr);

    while (!g_finish) {

        map<int, string> pid_names;
        int ret = get_pid_names(argv[1], pid_names);
        if (ret != 0) {
            return -1;
        }

        map<int, string> add_pid_names, del_pid_names;
        diff_pid_names(pid_names, cur_pid_names, add_pid_names, del_pid_names);
/*
        pid_desc(pid_names, "curr pids");
        pid_desc(cur_pid_names, "last pids");
*/
        monitor_desc(add_pid_names, "add monitor");
        monitor_desc(del_pid_names, "del monitor");

        for (const auto& item : add_pid_names) {
            auto ptr = shared_ptr<PidInfo>(new PidInfo{item.first, item.second, false});
            pid_info_dt[item.first] = ptr;
            thread([ptr] {
                pid_monitor(ptr);
            }).detach();
        }

        vector<int> del_pids;
        for (const auto& item : del_pid_names) {

            auto iter = pid_info_dt.find(item.first);
            if (iter == pid_info_dt.end()) {
                continue ;
            }

            auto ptr = iter->second;
            ptr->finish = true;

            dump_files("process exit", {item.first});

            pid_info_dt.erase(iter);
        }

        cur_pid_names = pid_names;

        sleep(1);
    }

    dump_files("program exit");

    return 0;
}

